import numpy as np
import random
import time
import matplotlib.pyplot as plt

# Generate a random matrix (values between 0 and 100)
def generate_matrix(n):
    return np.random.randint(0, 101, size=(n, n))

# Generate a random vector (values between 0 and 100)
def generate_vector(n):
    return np.random.randint(0, 101, size=(n, 1))

# SOR iteration
def solve_SOR(A, b, x0, K, e, w):
    n = len(A)
    x0 = x0.flatten()
    k = 0  # Iteration count
    err = np.inf  # Initialize error as infinity

    while k < K and err > e:
        k += 1
        x1 = x0.copy()
        
        for i in range(n):
            sum1 = np.dot(A[i, :i], x0[:i])    # Calculate upper triangular part
            sum2 = np.dot(A[i, i+1:], x1[i+1:])  # Calculate lower triangular part
            x0[i] = (1 - w) * x1[i] + (w / A[i][i]) * (b[i] - sum1 - sum2)

        # Calculate error (maximum norm between k+1 and k)
        err = np.max(np.abs(x1 - x0))

    return k, x0.reshape(-1, 1)

if __name__ == "__main__":
    # Define matrix A and vector b
    A = np.array([
        [31, -13, 0, 0, 0, -10, 0, 0, 0],
        [-13, 35, -9, 0, -11, 0, 0, 0, 0],
        [0, -9, 31, -10, 0, 0, 0, 0, 0],
        [0, 0, -10, 79, -30, 0, 0, 0, -9],
        [0, 0, 0, -30, 57, -7, 0, -5, 0],
        [0, 0, 0, 0, -7, 47, -30, 0, 0],
        [0, 0, 0, 0, 0, -30, 41, 0, 0],
        [0, 0, 0, 0, -5, 0, 0, 27, -2],
        [0, 0, 0, -9, 0, 0, 0, -2, 29],
    ])
    b = np.array([[-15], [27], [-23], [0], [-20], [12], [-7], [7], [10]])

    # Initial guess vector
    x0 = np.zeros((9, 1))
    # Maximum number of iterations
    K = 100
    # Error tolerance
    e = 1e-8
    # Relaxation factor
    w = 1.18 # 迭代得到的最优解

    # Timing
    start_time = time.time()
    
    # SOR iteration
    k, x = solve_SOR(A, b, x0, K, e, w)
    end_time = time.time()
    print("SOR Iteration")
    print("Number of iterations:", k)
    print("Approximate solution:\n", x)
    print("Elapsed time: {:.4f} seconds".format(end_time - start_time))

    # Gauss-Seidel iteration (using SOR with w=1)
    x0 = np.zeros((9, 1))
    k_gs, x_gs = solve_SOR(A, b, x0, K, e, 1)
    print("Gauss-Seidel Iteration")
    print("Number of iterations:", k_gs)
    print("Approximate solution:\n", x_gs)

    # Relaxation factor and iteration count relationship
    k_list = []
    w_list = []
    
    for i in range(1, 100):
        w = i / 50  # 松弛因子
        x0 = np.zeros((9, 1))
        k, _ = solve_SOR(A, b, x0, K, e, w)
        k_list.append(k)
        w_list.append(w)

    # Best relaxation factor
    w_best = w_list[k_list.index(min(k_list))]
    print("Best relaxation factor:", w_best)

    # Plotting the relationship between iteration count and relaxation factor
    plt.plot(w_list, k_list, marker='o')
    plt.title("Relationship between Relaxation Factor and Number of Iterations")
    plt.xlabel("Relaxation Factor w")
    plt.ylabel("Number of Iterations k")
    plt.grid(True)
    plt.show()

    # Print relaxation factors and their corresponding iteration counts
    for i in range(len(w_list)):
        print("Relaxation factor: {:.2f}, Number of iterations: {}".format(w_list[i], k_list[i]))

    # 进行解验证
    b_hat = np.dot(A, x)
    assert np.allclose(b, b_hat, atol=1e-7), "解1验证失败：Ax ≠ b"

    b_hat_2 = np.dot(A, x_gs)

    # 对比验证
    assert np.allclose(b, b_hat_2, atol=1e-7), "解2验证失败：Ax ≠ b"